package com.example.javapoet.demo.advanced;

import com.squareup.javapoet.*;

import javax.lang.model.element.Modifier;
import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.List;

/**
 * JavaPoet高级示例 - 生成访问者模式代码
 */
public class VisitorPatternGeneration {

    public static void main(String[] args) {
        generateVisitorPattern();
    }

    /**
     * 生成访问者模式的代码
     */
    public static void generateVisitorPattern() {
        // 定义包名
        String basePackage = "com.example.visitor";
        
        // 1. 创建Visitor接口
        TypeSpec visitorInterface = TypeSpec.interfaceBuilder("ShapeVisitor")
                .addModifiers(Modifier.PUBLIC)
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Circle"), "circle")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Rectangle"), "rectangle")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Triangle"), "triangle")
                        .build())
                .build();
        
        // 2. 创建Shape接口
        TypeSpec shapeInterface = TypeSpec.interfaceBuilder("Shape")
                .addModifiers(Modifier.PUBLIC)
                .addMethod(MethodSpec.methodBuilder("accept")
                        .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "ShapeVisitor"), "visitor")
                        .build())
                .build();
        
        // 3. 创建具体形状类
        // 3.1 Circle类
        TypeSpec circleClass = TypeSpec.classBuilder("Circle")
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ClassName.get(basePackage, "Shape"))
                .addField(FieldSpec.builder(double.class, "radius")
                        .addModifiers(Modifier.PRIVATE)
                        .build())
                .addMethod(MethodSpec.constructorBuilder()
                        .addModifiers(Modifier.PUBLIC)
                        .addParameter(double.class, "radius")
                        .addStatement("this.radius = radius")
                        .build())
                .addMethod(MethodSpec.methodBuilder("getRadius")
                        .addModifiers(Modifier.PUBLIC)
                        .returns(double.class)
                        .addStatement("return radius")
                        .build())
                .addMethod(MethodSpec.methodBuilder("accept")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "ShapeVisitor"), "visitor")
                        .addStatement("visitor.visit(this)")
                        .build())
                .build();
        
        // 3.2 Rectangle类
        TypeSpec rectangleClass = TypeSpec.classBuilder("Rectangle")
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ClassName.get(basePackage, "Shape"))
                .addField(FieldSpec.builder(double.class, "width")
                        .addModifiers(Modifier.PRIVATE)
                        .build())
                .addField(FieldSpec.builder(double.class, "height")
                        .addModifiers(Modifier.PRIVATE)
                        .build())
                .addMethod(MethodSpec.constructorBuilder()
                        .addModifiers(Modifier.PUBLIC)
                        .addParameter(double.class, "width")
                        .addParameter(double.class, "height")
                        .addStatement("this.width = width")
                        .addStatement("this.height = height")
                        .build())
                .addMethod(MethodSpec.methodBuilder("getWidth")
                        .addModifiers(Modifier.PUBLIC)
                        .returns(double.class)
                        .addStatement("return width")
                        .build())
                .addMethod(MethodSpec.methodBuilder("getHeight")
                        .addModifiers(Modifier.PUBLIC)
                        .returns(double.class)
                        .addStatement("return height")
                        .build())
                .addMethod(MethodSpec.methodBuilder("accept")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "ShapeVisitor"), "visitor")
                        .addStatement("visitor.visit(this)")
                        .build())
                .build();
        
        // 3.3 Triangle类
        TypeSpec triangleClass = TypeSpec.classBuilder("Triangle")
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ClassName.get(basePackage, "Shape"))
                .addField(FieldSpec.builder(double.class, "base")
                        .addModifiers(Modifier.PRIVATE)
                        .build())
                .addField(FieldSpec.builder(double.class, "height")
                        .addModifiers(Modifier.PRIVATE)
                        .build())
                .addMethod(MethodSpec.constructorBuilder()
                        .addModifiers(Modifier.PUBLIC)
                        .addParameter(double.class, "base")
                        .addParameter(double.class, "height")
                        .addStatement("this.base = base")
                        .addStatement("this.height = height")
                        .build())
                .addMethod(MethodSpec.methodBuilder("getBase")
                        .addModifiers(Modifier.PUBLIC)
                        .returns(double.class)
                        .addStatement("return base")
                        .build())
                .addMethod(MethodSpec.methodBuilder("getHeight")
                        .addModifiers(Modifier.PUBLIC)
                        .returns(double.class)
                        .addStatement("return height")
                        .build())
                .addMethod(MethodSpec.methodBuilder("accept")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "ShapeVisitor"), "visitor")
                        .addStatement("visitor.visit(this)")
                        .build())
                .build();
        
        // 4. 创建具体访问者类
        // 4.1 面积计算访问者
        TypeSpec areaVisitorClass = TypeSpec.classBuilder("AreaCalculator")
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ClassName.get(basePackage, "ShapeVisitor"))
                .addField(FieldSpec.builder(double.class, "totalArea")
                        .addModifiers(Modifier.PRIVATE)
                        .initializer("0.0")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Circle"), "circle")
                        .addStatement("totalArea += Math.PI * circle.getRadius() * circle.getRadius()")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Rectangle"), "rectangle")
                        .addStatement("totalArea += rectangle.getWidth() * rectangle.getHeight()")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Triangle"), "triangle")
                        .addStatement("totalArea += 0.5 * triangle.getBase() * triangle.getHeight()")
                        .build())
                .addMethod(MethodSpec.methodBuilder("getTotalArea")
                        .addModifiers(Modifier.PUBLIC)
                        .returns(double.class)
                        .addStatement("return totalArea")
                        .build())
                .build();
        
        // 4.2 绘图访问者
        TypeSpec drawVisitorClass = TypeSpec.classBuilder("DrawVisitor")
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ClassName.get(basePackage, "ShapeVisitor"))
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Circle"), "circle")
                        .addStatement("$T.out.println($S + circle.getRadius())", System.class, "绘制圆形，半径: ")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Rectangle"), "rectangle")
                        .addStatement("$T.out.println($S + rectangle.getWidth() + $S + rectangle.getHeight())", 
                                System.class, "绘制矩形，宽度: ", ", 高度: ")
                        .build())
                .addMethod(MethodSpec.methodBuilder("visit")
                        .addModifiers(Modifier.PUBLIC)
                        .addAnnotation(Override.class)
                        .returns(void.class)
                        .addParameter(ClassName.get(basePackage, "Triangle"), "triangle")
                        .addStatement("$T.out.println($S + triangle.getBase() + $S + triangle.getHeight())", 
                                System.class, "绘制三角形，底边: ", ", 高度: ")
                        .build())
                .build();
        
        // 5. 创建客户端示例类
        TypeSpec clientClass = TypeSpec.classBuilder("VisitorPatternDemo")
                .addModifiers(Modifier.PUBLIC)
                .addMethod(MethodSpec.methodBuilder("main")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .returns(void.class)
                        .addParameter(String[].class, "args")
                        .addStatement("$T<$T> shapes = $T.asList(\n" +
                                "    new $T(5.0),\n" +
                                "    new $T(4.0, 6.0),\n" +
                                "    new $T(3.0, 4.0)\n" +
                                ")",
                                List.class, ClassName.get(basePackage, "Shape"), Arrays.class,
                                ClassName.get(basePackage, "Circle"),
                                ClassName.get(basePackage, "Rectangle"),
                                ClassName.get(basePackage, "Triangle"))
                        .addCode("\n")
                        .addComment("使用面积计算访问者")
                        .addStatement("$T areaCalculator = new $T()", 
                                ClassName.get(basePackage, "AreaCalculator"),
                                ClassName.get(basePackage, "AreaCalculator"))
                        .beginControlFlow("for ($T shape : shapes)", ClassName.get(basePackage, "Shape"))
                        .addStatement("shape.accept(areaCalculator)")
                        .endControlFlow()
                        .addStatement("$T.out.println($S + areaCalculator.getTotalArea())", 
                                System.class, "所有形状的总面积: ")
                        .addCode("\n")
                        .addComment("使用绘图访问者")
                        .addStatement("$T drawVisitor = new $T()", 
                                ClassName.get(basePackage, "DrawVisitor"),
                                ClassName.get(basePackage, "DrawVisitor"))
                        .beginControlFlow("for ($T shape : shapes)", ClassName.get(basePackage, "Shape"))
                        .addStatement("shape.accept(drawVisitor)")
                        .endControlFlow()
                        .build())
                .build();
        
        // 创建Java文件
        List<JavaFile> javaFiles = Arrays.asList(
                JavaFile.builder(basePackage, shapeInterface).build(),
                JavaFile.builder(basePackage, visitorInterface).build(),
                JavaFile.builder(basePackage, circleClass).build(),
                JavaFile.builder(basePackage, rectangleClass).build(),
                JavaFile.builder(basePackage, triangleClass).build(),
                JavaFile.builder(basePackage, areaVisitorClass).build(),
                JavaFile.builder(basePackage, drawVisitorClass).build(),
                JavaFile.builder(basePackage, clientClass).build()
        );

        try {
            // 输出到文件
            Path path = Paths.get("src/main/java");
            for (JavaFile javaFile : javaFiles) {
                javaFile.writeTo(new File(path.toUri()));
            }
            System.out.println("成功生成访问者模式相关文件");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}